import os,time,sys,pdb
from pylab import plot,title,show,figure,savefig
from random import random
from math import sqrt,cos,sin,pi,acos,atan
from mpi4py import MPI
from numpy import array,dot,cross,isnan,linspace,empty,savez,c_,ravel, mgrid
from v3dsolver.vector import Vec
from v3dsolver.fluid_elements.fluid_elems import Saffman,FreeStream,FluidElem
from v3dsolver.boundary_manager import BoundaryManager,Geometry,Body
from v3dsolver.tree_manager.distribution_manager import DistributionManager,get_domain
from v3dsolver.tree_manager.vel_calculator import VelCalculator
from v3dsolver.tree_manager.tree_manager import TreeManager
from v3dsolver.advection_manager import AdvectionManager
from v3dsolver.stretching_manager import StretchingManager
from v3dsolver.diffusion_manager import DiffusionManager
from v3dsolver.merging_manager import MergingManager
from v3dsolver.annihilation_manager import AnnihilationManager
from v3dsolver.vorticity_interpolator import VorticityInterpolator
from enthought.tvtk.api import tvtk,write_data
from v3dsolver.writer import write

class Solver:
    """
    The overall 3D vortex solver, which combines functions like releasing blobs 
    from the boundary, advection, stretching, diffusion, merging, eating and 
    reflecting from boundary to simulate 3D flow past arbitrarily shaped
    solid bodies.
    """
    def __init__(self,bdry_file,rawdata_dir,blob_rad,delta,npmax,npseudo,dt,
                 freestream_val,dx_stretching,nu,R_bin,merge_frac,
                 cutoff_fraction):
        """
        Initialization of Solver. It initializes various parameters used in the
        simulation.
        
        bdry_file: filename of the vtk file representing the boundary
        rawdata_dir: directory path (relative to project directory) 
                     where simulation output will be stored in raw format
        blob_rad: radius of vortex blobs
        delta: distance of released blob from the boundary surface
        npmax: maximum no. of particles in a cell for it to be childless
        npseudo: no. of pseudo particles to be generated per cell
        dt: timestep of the simulation
        freestream_val: value of freestream velocity
        dx_stretching: distance between points on the same stretch stensil
        nu: kinematic viscosity
        R_bin: radius of binning (keep around 5*sqrt(nu*dt))
        merge_frac: merging radius as a fraction of cellsize = R_bin*sqrt(nu*dt)
        cutoff_fraction: fraction of avg vorticity below which a particle will
                         be discarded
        """
        # set all parameters
        self.namespace = {'fs':fs,'rad':rad,'npmax':npmax,
                          'npseudo':npseudo,'delta':delta,'dt':dt,'dx':dx,
                          'nu':nu,'R':R,'merge_frac':merge_frac,
                          'cutoff_frac':cutoff_frac,'mode':mode}
        self.params = {'b':['rad','fs','npmax','npseudo','delta','mode'],
                       'a':['rad','fs','npmax','npseudo','dt','mode'],
                       's':['dx','dt','fs','npmax','npseudo','rad','mode'],
                       'd':['nu','dt','R','mode'],
                       'm':['merge_frac'],
                       'n':['cutoff_frac'],
                       'r':['nu','dt','R']}
        self.components = components
        for component in self.components:
            for param in self.params[component]:
                self.__setattr__(param,self.params[param])
        self.rad = rad # universal parameter
        
        
        self.geometry = Geometry([Body(bdry_file,blob_rad,delta)]) 
        self.delta = delta
        self.blob_rad = blob_rad
        self.npmax = npmax
        self.npseudo = npseudo
        self.dt = dt
        self.fs_val = freestream_val
        self.dx = dx_stretching
        self.nu = nu
        self.R_bin = R_bin
        self.merge_frac = merge_frac
        self.cutoff_frac = cutoff_fraction
        self.gridsize = 25
        self.gridpts, self.insidegridpts,self.totalgridpts = self.get_grid_pts2(self.gridsize)
    
    def initialize(self):
        """
        Initializes flow simulation by creation of boundary using geometry data 
        and/or by generation of initial vortex particles using outside 
        blob data. 
        """
        # initialize bounded flows
        if self.geometry:
            self.boundary_manager = BoundaryManager(self.geometry,self.rad,self.fs,self.mode)
            self.advection_manager = AdvectionManager(self.dt,self.fs,self.npmax,self.npseudo,self.rad,self.mode)
            self.stretching_manager = StretchingManager(self.dx,self.dt,self.fs,self.npmax,self.npseudo,self.rad,self.mode)
            self.diffusion_manager = DiffusionManager(self.nu,self.dt,self.R_bin,self.merge_frac,self.mode,self.geometry)
            self.annihilation_manager = AnnihilationManager(self.cutoff_frac)
        
        # initialize unbounded flows
        self.flowblobs = self.get_elems()
        self.bblobs = []
        
    def boundary(self):
        """
        Solves no-penetration and no-slip boundary conditions on the 
        boundary.
        """
        if self.geometry:
            rblobs,self.bblobs = self.boundary_manager.get_released_particles(self.flowblobs)
            self.flowblobs += rblobs
            
    def advect(self):
        """
        Computes velocity field due to existing blobs and moves all the
        particles according to their calculated velocities.
        """
        self.advection_manager.move_Euler(self.flowblobs+self.bblobs)
        
    def stretch(self):
        """
        Computes gradient of velocity field using ghost particles around
        each particle and uses it to update vorticity according to the
        stretching equation.
        """
        self.stretching_manager.update_vors(self.flowblobs+self.bblobs)
        
    def diffuse(self):
        """
        Diffuses particles' vorticities to their neighbour particles. 
        Adds new particles if sufficient neighbours are not present. 
        """
        self.diffusion_manager.diffuse(self.flowblobs)
        
    def merge(self):
        """
        Merges the particles if they are too close to each other.
        """
        self.diffusion_manager.merge(self.flowblobs)
        
    def reflect(self):
        """
        Finds neighbours to the panels and reflects any flow blobs going
        in.
        """
        self.diffusion_manager.reflect()
        
    def annihilate(self):
        """
        Annihilates particles which have significantly small vorticity.
        """
        self.annihilation_manager.annihilate()
        
    def write(self):
        """
        Writes obtained simulation data to files.
        """
        self.writer.write()
    
    def solve(self,tf,start=0):
        """
        Solves flow till time tf.

        tf: time uptil which simulation needs to be run
        start: if non-zero, it will pick up data from the file corresponding
               to that iteration and continue from there.
        """
        self.tf = tf
        self.start = start
        self.t = start*self.dt
        
        self.initialize()
        
        while self.t < self.tf:
            for module in self.modules.values():
                module()
            self.t += self.dt

if __name__=="__main__":
    casename = sys.argv[1]
    vortexsolver = Solver(bdry_file=casename+'.vtk',rawdata_dir='examples/'+casename+'/raw/Re100/',
                          blob_rad=0.1,delta=0.1,npmax=5,npseudo=16,dt=0.01,
                          freestream_val=Vec(0.,0.,10.),
                          dx_stretching=0.001,nu=9.e-2,R_bin=5,
                          merge_frac=0.2,cutoff_fraction=0.005)
    vortexsolver.solve(tf=0.5,components='badnw')
